import json
import os
from datetime import timedelta
from typing import Dict, Union
from uuid import uuid1

from celery import Celery, Task
from flask import g, has_app_context, has_request_context, request
from redis import Redis

from configs.settings import celery_config
import logging


logger = logging.getLogger(__name__)
# from celery.contrib import rdb
# rdb.set_trace()


class ContextTask(Task):

    DUMP_KEY = "TASK_DUMP_REQUEST_CONTEXT::"

    def build_request_context(self, user_id: int, role: str):
        """构建必要请求上下文信息，为部分模块调用时提供身份信息"""
        redis = self.app.flask_app.extensions["redis"]
        context_info = {
            "method": "POST",
            "headers": {},
            "role": role,
            "user_id": user_id,
        }
        redis.set(
            f"{self.DUMP_KEY}{self.request.id}", json.dumps(context_info)
        )

    def dump_current_request_context(self, task_id, login_info=None) -> Dict:
        """转储flask request必要上下文信息"""

        redis: Redis = self.app.flask_app.extensions["redis"]
        db = self.app.flask_app.extensions["sqlalchemy"]  # noqa

        cache_context = redis.get(f"{self.DUMP_KEY}{task_id}")
        if cache_context:
            return json.loads(cache_context.decode())

        if not has_request_context():
            return None

        user = getattr(g, "flask_httpauth_user", None)
        context_info = {
            "method": request.method,
            "headers": dict(request.headers),
            "role": request.cookies.get("role", None),
            "user_id": user.id if user else None,
        }

        redis.set(f"{self.DUMP_KEY}{task_id}", json.dumps(context_info))

        return context_info

    def expand_rest(self, rest: dict, login_info=None):
        task_id = self.request.id or uuid1().hex
        rest.setdefault("headers", {})
        # 通过celery发送到任务后，可在self.request.flask_request_context中获取
        if not login_info:
            rest["headers"]["flask_request_context"] = self.dump_current_request_context(
                task_id
            )
        else:
            if not isinstance(login_info, dict):
                login_info = {}
            login_info.setdefault("method", "GET")
            login_info.setdefault("headers", {})
            login_info.setdefault("role", None)
            login_info.setdefault("user_id", None)
            rest["headers"]["flask_request_context"] = login_info
        rest["task_id"] = task_id

    def get_login_info(self):
        return self.request.flask_request_context

    def get_login_user(self, login_role="teacher"):
        from app.models import User

        flask_request_context = self.request.flask_request_context

        if not flask_request_context:
            user = User.query.filter(User.username == "admin").first()
            role = login_role
            return user, role

        user_id = flask_request_context.get("user_id", None)
        if not user_id:
            user = User.query.filter(User.username == "admin").first()
        else:
            user = User.query.get(user_id)

        role = flask_request_context.get("role", login_role)

        return user, role

    def delay(self, *args, login_info=None, **kwargs):
        return self.apply_async(args, kwargs, login_info=login_info)

    def apply_async(self, args=None, kwargs=None, login_info=None, **rest):
        self.expand_rest(rest, login_info=login_info)
        return super().apply_async(args, kwargs, **rest)

    def apply(self, args=None, kwargs=None, login_info=None, **rest):
        self.expand_rest(rest, login_info=login_info)
        return super().apply(args, kwargs, **rest)

    def retry(self, args=None, kwargs=None, login_info=None, **rest):
        self.expand_rest(rest, login_info=login_info)
        return super().retry(args, kwargs, **rest)

    def __call__(self, *args, **kwargs):

        task_id = self.request.id

        # 无需重复加载flask上下文
        if has_app_context():
            os.environ["FLASK_CONTEXT_IN_CELERY"] = "true"
            g.task_id = task_id
            self.app.flask_app.extensions["sqlalchemy"].db.session.close()
            result = super(ContextTask, self).__call__(*args, **kwargs)
        else:
            with self.app.flask_app.app_context():
                os.environ["FLASK_CONTEXT_IN_CELERY"] = "true"
                g.task_id = task_id
                self.app.flask_app.extensions["sqlalchemy"].db.session.close()
                result = super(ContextTask, self).__call__(*args, **kwargs)
        return result

    def after_return(self, status, retval, task_id, args, kwargs, einfo):
        redis: Redis = self.app.flask_app.extensions["redis"]
        redis.delete(f"{self.DUMP_KEY}{task_id}")


def create_celery_app(flask_app, config=None):
    celery = Celery(flask_app.import_name or celery_config.celery_name or "default")

    if config is None:
        config = celery_config

    config.imports = ["app.api", "tasks.tasks"]
    celery.config_from_object(config, force=True)

    celery.conf.beat_schedule = {
        "check_active_vm_plc_status_on_every_second": {
            "task": "tasks.tasks.check_all_range_plc_status",  # noqa
            "schedule": timedelta(seconds=2),
        },
        "verify_for_system_run_conditions_on_every_day": {
            "task": "tasks.tasks.verify_for_system_run_conditions",  # noqa
            "schedule": timedelta(days=1),
        },
    }
    celery.autodiscover_tasks()

    celery.Task = ContextTask

    if not hasattr(celery, "flask_app"):
        celery.flask_app = flask_app

    return celery


class CeleryExtension:
    def __init__(self, app=None, create_celery=None):
        self.create_celery_app = create_celery or create_celery_app
        self._celery = None
        if app is not None:
            self.init_app(app)

    def init_app(self, app, config=None):
        if not hasattr(app, "extensions"):
            app.extensions = {}
        app.extensions["celery_ext"] = self
        # celery app
        self._celery = self.create_celery_app(app, config)

    def __getattr__(self, name):
        return getattr(self._celery, name)

    def __getitem__(self, name):
        return self._celery[name]

    def __setitem__(self, name, value):
        self._celery[name] = value

    def __delitem__(self, name):
        del self._celery[name]

    @property
    def channel_client(self: Celery):
        """celery后端客户端"""
        return self.connection().channel().client

    @property
    def inspect(self: Celery):
        """task控制器"""
        return self.control.inspect()


celery_app: Union[Celery, CeleryExtension] = CeleryExtension()
